'''
	Copyright (C) 2010  Pier Maria Mattioli
	
	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.
	
	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.
	
	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
	
'''

from clusty import Manager, Cluster
import unittest
import random
import threading
import time

chars = list('1234567890!@#$%^&*()-_+=qwertyuiop[]{}asdfghjkl;:"|`~zxcvbnm,<.>/?')

class ClusterTestCase(unittest.TestCase):
	'''
	Methods to test here:
	 |  catch_lock(self, lock)
	 |  
	 |  delete(self, key)
	 |  
	 |  get(self, key, default=None)
	 |  
	 |  get_lock(self, lock, times=0, delay=0.0050000000000000001)
	 |  
	 |  keys(self, pattern)
	 |  
	 |  mget(self, mapping)
	 |  
	 |  mset(self, mapping)
	 |  
	 |  release_lock(self, lock)
	 |  
	 |  set(self, key, value)
	 |  
	 |  setnx(self, key, value)
	'''

	def setUp(self):
		self.nodes = ['localhost:6379:0', 'localhost:6379:1']
		self.redundancy = 2
		self.cluster = Cluster(self.nodes, self.redundancy)
		for node in self.cluster.container.values():
			node.flushdb()
	
	def tearDown(self):
		for node in self.cluster.container.values():
			node.flushdb()
		self.cluster.shutdown()
	
	def test_1_catch_lock(self):
		self.cluster.release_lock('lock.test')
		self.assertTrue(self.cluster.catch_lock('lock.test'))
		self.assertFalse(self.cluster.catch_lock('lock.test'))
		self.cluster.release_lock('lock.test')
		self.assertTrue(self.cluster.catch_lock('lock.test'))
		self.assertFalse(self.cluster.catch_lock('lock.test'))
		self.cluster.release_lock('lock.test')
	
	def test_2_get_lock(self):
		self.cluster.release_lock('lock.test')
		self.assertTrue(self.cluster.get_lock('lock.test', times=0, delay=0.0050000000000000001))
		self.assertFalse(self.cluster.get_lock('lock.test', times=5, delay=0.0050000000000000001))
		self.cluster.release_lock('lock.test')
		
	def test_3_single_operations(self):
		self.cluster.set('testing', 'ohhhhhh yesssssss!')
		self.assertEquals(self.cluster.get('testing'), 'ohhhhhh yesssssss!')
		self.cluster.delete('testing')
		self.assertEquals(self.cluster.get('testing'), None)
		self.assertEquals(self.cluster.get('testing', False), False)
		self.cluster.setnx('testing', 'ohhhhhh yesssssss!')
		self.cluster.setnx('testing', 'ohhhhhh nooooooooo!')
		self.assertEquals(self.cluster.get('testing'), 'ohhhhhh yesssssss!')
		self.cluster.delete('testing')
		
	def test_4_multi_operations(self):
		self.assertEquals(self.cluster.keys('*'), ['ctr::nodes::all'])
		# 
		self.mydict = {}
		for i in range(1000):
			random.shuffle(chars)
			key = ''.join(random.sample(chars, 12))
			value = ''.join(random.sample(chars, 24))
			self.mydict[key] = value
		self.cluster.mset(self.mydict.items())
		# 
		k1 = self.mydict.keys()
		k1.sort()
		k2 = self.cluster.keys('*')
		k2 = [k for k in k2 if k != 'ctr::nodes::all']
		k2.sort()
		self.assertEquals(k1, k2)
		# 
		v1 = [self.mydict[k] for k in k1] # [:10]]
		v2 = self.cluster.mget(k1) # [:10])
		self.assertEquals(v1, v2)
		# 
		for k in k1:
			self.cluster.delete(k)
		self.assertEquals(self.cluster.keys('*'), ['ctr::nodes::all'])
		
	
class ManagerTestCase(unittest.TestCase):
	'''
	Methods to test here:
     |  add_node(self, node_id)
     |  
     |  change_redundancy(self, new_value)
     |  
     |  del_node(self, node_id)
     |  
     |  init_node(self, node)
     |  
     |  rebalance(self, node_id=None)
	'''

	def setUp(self):
		self.nodes = ['localhost:6379:0', 'localhost:6379:1']
		self.redundancy = 2
		self.manager = Manager(self.nodes, self.redundancy)
		for node in self.manager.container.values():
			node.flushdb()
		self.mydict = {}
		for i in range(1000):
			random.shuffle(chars)
			key = ''.join(random.sample(chars, 12))
			value = ''.join(random.sample(chars, 24))
			self.mydict[key] = value
		self.manager.mset(self.mydict.items())
		self.n_keys = len(self.mydict.keys())

	def tearDown(self):
		for node in self.manager.container.values():
			node.flushdb()
		self.manager.shutdown()
		
	def test_1_add_node(self):

		def _count_keys():
			self.manager._check_container()
			count = 0
			for node in self.manager.container.values():
				keys = self.manager.nodes.do_keys(node, '*')
				count += len(keys)
			return count			
				
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 2)
		
		self.manager.add_node('localhost:6379:2')
		self.assertEquals(len(self.manager.container.keys()), 3)
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 3)
		self.manager.rebalance('localhost:6379:2')
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 3)
		
	def test_2_del_node(self):

		def _count_keys():
			self.manager._check_container()
			count = 0
			for node in self.manager.container.values():
				keys = self.manager.nodes.do_keys(node, '*')
				count += len(keys)
			return count			

		self.manager.del_node('localhost:6379:2')
		self.assertEquals(len(self.manager.container.keys()), 2)
		self.manager.rebalance('localhost:6379:2')
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 2)

	def test_3_change_redundancy(self):

		def _count_keys():
			self.manager._check_container()
			count = 0
			for node in self.manager.container.values():
				keys = self.manager.nodes.do_keys(node, '*')
				count += len(keys)
			return count			
		
		self.manager.add_node('localhost:6379:2')
		self.manager.rebalance('localhost:6379:2')
		self.manager.change_redundancy(3)
		self.manager.rebalance()
		self.assertEquals(_count_keys(), (self.n_keys * 3) + 3)

		self.manager.change_redundancy(2)
		self.manager.rebalance()
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 3)

		
class HotOpTestCase(unittest.TestCase):
	def setUp(self):
		self.nodes = ['localhost:6379:0', 'localhost:6379:1']
		self.redundancy = 2
		self.cluster = Cluster(self.nodes, self.redundancy)
		self.manager = Manager(self.nodes, self.redundancy)		
		self.mydict = {}
		for i in range(1000):
			random.shuffle(chars)
			key = ''.join(random.sample(chars, 12))
			value = ''.join(random.sample(chars, 24))
			self.mydict[key] = value
		self.n_keys = len(self.mydict.keys())
	
	def tearDown(self):
		for node in self.cluster.container.values():
			node.flushdb()
		self.manager.shutdown()
		self.cluster.shutdown()
		
	def test_1_single_operations(self):
		def _count_keys():
			self.manager._check_container()
			count = 0
			for node in self.manager.container.values():
				keys = self.manager.nodes.do_keys(node, '*')
				count += len(keys)
			return count	
		
		def bulk_set():
			for k,v in self.mydict.iteritems():
				self.cluster.set(k,v)

		def bulk_delete():
			for k,v in self.mydict.iteritems():
				self.cluster.delete(k)

		t = threading.Thread(target=bulk_set)
		t.start()
		self.manager.add_node('localhost:6379:2')
		self.manager.rebalance('localhost:6379:2')
		t.join()
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 3)

		t = threading.Thread(target=bulk_delete)
		t.start()
		self.manager.add_node('localhost:6379:3')
		self.manager.rebalance('localhost:6379:3')
		t.join()
		self.assertEquals(_count_keys(), 4)

		t = threading.Thread(target=bulk_set)
		t.start()
		self.manager.del_node('localhost:6379:2')
		self.manager.rebalance('localhost:6379:2')
		t.join()
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 3)

		t = threading.Thread(target=bulk_delete)
		t.start()
		self.manager.del_node('localhost:6379:3')
		self.manager.rebalance('localhost:6379:3')
		t.join()
		self.assertEquals(_count_keys(), 2)


	def test_2_multi_operations(self):
		def _count_keys():
			self.manager._check_container()
			count = 0
			for node in self.manager.container.values():
				keys = self.manager.nodes.do_keys(node, '*')
				count += len(keys)
			return count	
		
		def bulk_mset():
			self.cluster.mset(self.mydict.items())

		def bulk_delete():
			for k,v in self.mydict.iteritems():
				self.cluster.delete(k)

		t = threading.Thread(target=bulk_mset)
		t.start()
		self.manager.add_node('localhost:6379:2')				
		self.manager.rebalance('localhost:6379:2')
		t.join()
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 3)
		
		t = threading.Thread(target=bulk_delete)
		t.start()
		self.manager.add_node('localhost:6379:3')
		self.manager.rebalance('localhost:6379:3')
		t.join()
		self.assertEquals(_count_keys(), 4)

		t = threading.Thread(target=bulk_mset)
		t.start()
		self.manager.del_node('localhost:6379:2')
		self.manager.rebalance('localhost:6379:2')
		t.join()
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 3)

		t = threading.Thread(target=bulk_delete)
		t.start()
		self.manager.del_node('localhost:6379:3')
		self.manager.rebalance('localhost:6379:3')
		t.join()
		self.assertEquals(_count_keys(), 2)
	

	def test_3_locking_operations(self):
		def _count_keys():
			self.manager._check_container()
			count = 0
			for node in self.manager.container.values():
				keys = self.manager.nodes.do_keys(node, '*')
				count += len(keys)
			return count	
		
		def bulk_lock():
			for k,v in self.mydict.iteritems():
				self.cluster.get_lock(k)

		def bulk_unlock():
			for k,v in self.mydict.iteritems():
				self.cluster.release_lock(k)

		t = threading.Thread(target=bulk_lock)
		t.start()
		self.manager.add_node('localhost:6379:2')
		self.manager.rebalance('localhost:6379:2')
		t.join()
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 3)

		t = threading.Thread(target=bulk_unlock)
		t.start()
		self.manager.add_node('localhost:6379:3')
		self.manager.rebalance('localhost:6379:3')
		t.join()
		self.assertEquals(_count_keys(), 4)

		t = threading.Thread(target=bulk_lock)
		t.start()
		self.manager.del_node('localhost:6379:2')
		self.manager.rebalance('localhost:6379:2')
		t.join()
		self.assertEquals(_count_keys(), (self.n_keys * 2) + 3)

		t = threading.Thread(target=bulk_unlock)
		t.start()
		self.manager.del_node('localhost:6379:3')
		self.manager.rebalance('localhost:6379:3')
		t.join()
		self.assertEquals(_count_keys(), 2)		
		